# encoding: utf-8
"""
!/usr/bin/python3
@Author: Gao Shuo
@Time: 2019/3/29 14:31 
@ReadMe:
    Input: 
    Output: 
"""

from mesa.time import RandomActivation
# The below is needed for both notebooks and scripts
import matplotlib.pyplot as plt
from tutorial_model import *
from mesa.space import MultiGrid
import numpy as np
from mesa.batchrunner import BatchRunner

# model = MoneyModel(10)
# for i in range(10):
#     model.step()

# For a jupyter notebook add the following line:
# %matplotlib inline
# agent_wealth = [a.wealth for a in model.schedule.agents]
# plt.hist(agent_wealth)

# all_wealth = []
# for j in range(100):
#     # Run the model
#     model = MoneyModel(10)
#     for i in range(10):
#         model.step()
#
#     # Store the results
#     for agent in model.schedule.agents:
#         all_wealth.append(agent.wealth)
# plt.hist(all_wealth, bins=range(max(all_wealth) + 1))
# plt.show()

# 移动并给钱
# model = MoneyModel(50, 10, 10)
# for i in range(20):
#     model.step()
# agent_counts = np.zeros((model.grid.width, model.grid.height))
# for cell in model.grid.coord_iter():
#     cell_content, x, y = cell
#     agent_count = len(cell_content)
#     agent_counts[x][y] = agent_count
# plt.imshow(agent_counts, interpolation='nearest')
# plt.colorbar()
# plt.show()
# 画基尼系数

# datacollector 画图
model = MoneyModel(50, 10, 10)
for i in range(100):
    model.step()
gini = model.datacollector.get_model_vars_dataframe()
gini.plot()
plt.show()
plt.close('all')

agent_wealth = model.datacollector.get_agent_vars_dataframe()
print(agent_wealth.head())

end_wealth = agent_wealth.xs(5, level="Step")["Wealth"]
end_wealth.hist(bins=range(agent_wealth.Wealth.max()+1))

one_agent_wealth = agent_wealth.xs(14, level="AgentID")
one_agent_wealth.Wealth.plot()
plt.show()

#batch
# fixed_params = {"width": 10,
#                 "height": 10}
# variable_params = {"N": range(10, 100, 10)}
#
# batch_run = BatchRunner(MoneyModel,
#                         fixed_parameters=fixed_params,
#                         variable_parameters=variable_params,
#                         iterations=5,
#                         max_steps=10,
#                         model_reporters={"Gini": compute_gini})
# batch_run.run_all()
# run_data = batch_run.get_model_vars_dataframe()
# print("======run_data======")
# print(run_data.head())
# plt.scatter(run_data.N, run_data.Gini)
# plt.show()
# run_data.to_csv('test_output.csv')